from fastapi import APIRouter, File, UploadFile, BackgroundTasks, HTTPException, Query
from fastapi.responses import JSONResponse
import os
import shutil
from pathlib import Path
import json
from typing import List, Optional

import fitz  # PyMuPDF
from PIL import Image
import io
import pytesseract


from langchain_community.document_loaders import (
    PyPDFLoader, 
    TextLoader,
    UnstructuredWordDocumentLoader,
    UnstructuredFileLoader,
    UnstructuredPowerPointLoader,
    UnstructuredExcelLoader,
    UnstructuredCSVLoader
)
from langchain_core.documents import Document

from app.rag.retriever import Retriever
from app.core.logger import logger

# 将PDF_DIR重命名为UPLOADS_DIR并更改路径
UPLOADS_DIR = Path("uploads")
UPLOADS_DIR.mkdir(parents=True, exist_ok=True)

router = APIRouter()
retriever = Retriever()


def is_pdf_scanned(file_path: str) -> bool:
    """
    判断PDF是否为扫描件（通过检查PDF中是否包含文本内容）
    """
    try:
        doc = fitz.open(file_path)
        total_pages = doc.page_count
        text_content = ""
        
        # 检查前几页是否有文本内容
        pages_to_check = min(3, total_pages)  # 只检查前3页或总页数
        for page_num in range(pages_to_check):
            page = doc.load_page(page_num)
            text = page.get_text()
            text_content += text
            
        doc.close()
        
        # 如果文本内容很少，可能是扫描件
        return len(text_content.strip()) < 100
    except Exception as e:
        logger.error(f"Error checking if PDF is scanned: {e}")
        return False  # 出错时默认不是扫描件


def extract_text_from_scanned_pdf(file_path: str) -> str:
    """
    从扫描版PDF中提取文本（使用OCR）
    """
    try:
        doc = fitz.open(file_path)
        text = ""
        
        for page_num in range(doc.page_count):
            page = doc.load_page(page_num)
            pix = page.get_pixmap()
            
            # 将Pixmap转换为PIL Image
            img_data = pix.tobytes("ppm")
            img = Image.open(io.BytesIO(img_data))
            
            # 使用OCR提取文本
            page_text = pytesseract.image_to_string(img, lang='chi_sim+eng')
            text += f"--- Page {page_num + 1} ---\n{page_text}\n\n"
            
        doc.close()
        return text
    except Exception as e:
        logger.error(f"Error extracting text from scanned PDF: {e}")
        return ""


def save_document_file(file: UploadFile, save_path: Path) -> bool:
    """
    保存上传的文档文件
    """
    try:
        with open(save_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        return True
    except Exception as e:
        logger.error(f"Error saving file {file.filename}: {e}")
        return False


def load_document(file_path: Path, file_extension: str) -> list:
    """
    根据文件扩展名加载文档
    """
    try:
        if file_extension == ".pdf":
            # 检查是否为扫描件
            if is_pdf_scanned(str(file_path)):
                logger.info(f"PDF {file_path} is a scanned document, using OCR")
                text = extract_text_from_scanned_pdf(str(file_path))
                if text:
                    doc = Document(
                        page_content=text,
                        metadata={
                            "source": str(file_path),
                            "filename": file_path.name,
                            "filetype": "pdf_scanned"
                        }
                    )
                    return [doc]
                else:
                    logger.warning(f"Failed to extract text from scanned PDF {file_path}")
                    return []
            else:
                loader = PyPDFLoader(str(file_path))
                return loader.load()
        elif file_extension == ".txt":
            loader = TextLoader(str(file_path), encoding="utf-8")
            return loader.load()
        elif file_extension in [".doc", ".docx"]:
            loader = UnstructuredWordDocumentLoader(str(file_path))
            return loader.load()
        elif file_extension in [".md", ".markdown"]:
            # Markdown文件使用文本加载器
            loader = TextLoader(str(file_path), encoding="utf-8")
            return loader.load()
        elif file_extension in [".ppt", ".pptx"]:
            # PowerPoint文件
            loader = UnstructuredPowerPointLoader(str(file_path))
            return loader.load()
        elif file_extension in [".xls", ".xlsx"]:
            # Excel文件
            loader = UnstructuredExcelLoader(str(file_path))
            return loader.load()
        elif file_extension == ".csv":
            # CSV文件
            loader = UnstructuredCSVLoader(str(file_path))
            return loader.load()
        elif file_extension == ".ofd":
            # OFD文件（中国标准版式文件）
            logger.warning(f"OFD file support is not fully implemented: {file_path}")
            # 当前版本暂不支持OFD文件，返回空列表
            return []
        else:
            # 尝试通用加载器
            loader = UnstructuredFileLoader(str(file_path))
            return loader.load()
    except Exception as e:
        logger.error(f"Error loading document {file_path}: {e}")
        return []


@router.post("/knowledge/upload")
async def upload_document(file: UploadFile = File(...)):
    """
    上传并处理各种格式的文档文件
    支持 PDF（包括扫描件）、TXT、DOC、DOCX 等格式
    """
    # 支持的文件类型
    allowed_extensions = {
        ".pdf", ".txt", ".doc", ".docx",
        ".md", ".markdown",  # 新增Markdown
        ".ppt", ".pptx",     # 新增PowerPoint
        ".xls", ".xlsx",     # 新增Excel
        ".csv",              # 新增CSV
        ".ofd"               # OFD格式（中国标准版式文档格式）
    }
    
    # 检查文件扩展名
    file_extension = Path(file.filename).suffix.lower()
    if file_extension not in allowed_extensions:
        raise HTTPException(
            status_code=400, 
            detail=f"Unsupported file type: {file_extension}. "
                   f"Supported types: {', '.join(allowed_extensions)}"
        )

    # 确保保存目录存在
    upload_dir = UPLOADS_DIR
    upload_dir.mkdir(parents=True, exist_ok=True)
    
    # 保存文件
    file_path = upload_dir / file.filename
    if not save_document_file(file, file_path):
        raise HTTPException(status_code=500, detail="Failed to save file")

    # 加载文档
    documents = load_document(file_path, file_extension)
    if not documents:
        raise HTTPException(status_code=500, detail="Failed to load document content")

    # 分割文档
    from app.rag.splitter import split_documents
    split_docs = split_documents(documents)
    
    if not split_docs:
        raise HTTPException(status_code=500, detail="Failed to split document into chunks")

    # 添加到向量数据库
    try:
        retriever.vector_store.add_documents(split_docs)
        logger.info(f"Successfully processed {file.filename} with {len(split_docs)} chunks")
        return JSONResponse(
            status_code=200, 
            content={
                "code": 200, 
                "message": f"Upload successful, processed {len(split_docs)} chunks", 
                "filename": file.filename
            }
        )
    except Exception as e:
        logger.error(f"Failed to add documents to vector store: {e}")
        raise HTTPException(status_code=500, detail="Failed to process document")


@router.post("/knowledge/sync")
async def sync_knowledge():
    """手动触发同步"""
    try:
        retriever.ingest_from_pdf_dir()
        return {"code": 200, "message": "Sync completed"}
    except Exception as e:
        logger.error(f"Sync failed: {e}")
        raise HTTPException(status_code=500, detail="Sync failed")


@router.get("/knowledge/list")
async def list_documents(
    page: int = Query(1, ge=1, description="页码"),
    page_size: int = Query(10, ge=1, le=100, description="每页数量")
):
    """
    列出所有已上传的文档
    """
    try:
        # 确保上传目录存在
        upload_dir = UPLOADS_DIR
        upload_dir.mkdir(parents=True, exist_ok=True)
        
        # 获取所有文件
        files = list(upload_dir.iterdir())
        
        # 过滤出文件（排除目录）
        files = [f for f in files if f.is_file()]
        
        # 分页处理
        total = len(files)
        start_idx = (page - 1) * page_size
        end_idx = start_idx + page_size
        paginated_files = files[start_idx:end_idx]
        
        # 构造返回数据
        documents = []
        for file in paginated_files:
            stat = file.stat()
            documents.append({
                "filename": file.name,
                "size": stat.st_size,
                "modified_time": stat.st_mtime,
                "extension": file.suffix
            })
        
        return {
            "code": 200,
            "message": "success",
            "data": {
                "documents": documents,
                "total": total,
                "page": page,
                "page_size": page_size,
                "total_pages": (total + page_size - 1) // page_size
            }
        }
    except Exception as e:
        logger.error(f"Failed to list documents: {e}")
        raise HTTPException(status_code=500, detail="Failed to list documents")


@router.get("/knowledge/search")
async def search_documents(keyword: str = Query(..., description="搜索关键词")):
    """
    根据关键词搜索文档
    """
    try:
        # 确保上传目录存在
        upload_dir = UPLOADS_DIR
        upload_dir.mkdir(parents=True, exist_ok=True)
        
        # 获取所有文件
        files = list(upload_dir.iterdir())
        
        # 过滤出文件（排除目录）并根据关键词搜索
        matched_files = []
        for f in files:
            if f.is_file() and keyword.lower() in f.name.lower():
                stat = f.stat()
                matched_files.append({
                    "filename": f.name,
                    "size": stat.st_size,
                    "modified_time": stat.st_mtime,
                    "extension": f.suffix
                })
        
        return {
            "code": 200,
            "message": "success",
            "data": {
                "documents": matched_files,
                "total": len(matched_files)
            }
        }
    except Exception as e:
        logger.error(f"Failed to search documents: {e}")
        raise HTTPException(status_code=500, detail="Failed to search documents")


@router.delete("/knowledge/{filename}")
async def delete_document(filename: str):
    """
    删除指定文档
    """
    try:
        # 构造文件路径
        file_path = UPLOADS_DIR / filename
        
        # 检查文件是否存在
        if not file_path.exists():
            raise HTTPException(status_code=404, detail="Document not found")
        
        # 检查是否为文件（而不是目录）
        if not file_path.is_file():
            raise HTTPException(status_code=400, detail="Path is not a file")
        
        # 删除文件
        file_path.unlink()
        
        # TODO: 从向量数据库中删除相关文档（需要实现）
        # 这需要根据具体的向量存储实现来完成
        
        logger.info(f"Successfully deleted document: {filename}")
        return {
            "code": 200,
            "message": f"Document {filename} deleted successfully"
        }
    except HTTPException:
        # 重新抛出已知的HTTP异常
        raise
    except Exception as e:
        logger.error(f"Failed to delete document {filename}: {e}")
        raise HTTPException(status_code=500, detail="Failed to delete document")


@router.put("/knowledge/{filename}")
async def update_document(filename: str, file: UploadFile = File(...)):
    """
    更新指定文档（先删除再重新上传）
    """
    try:
        # 构造文件路径
        file_path = UPLOADS_DIR / filename
        
        # 检查文件是否存在
        if not file_path.exists():
            raise HTTPException(status_code=404, detail="Document not found")
        
        # 检查是否为文件（而不是目录）
        if not file_path.is_file():
            raise HTTPException(status_code=400, detail="Path is not a file")
        
        # 保存新文件（覆盖原文件）
        if not save_document_file(file, file_path):
            raise HTTPException(status_code=500, detail="Failed to save file")
        
        # 获取文件扩展名
        file_extension = Path(file.filename).suffix.lower() if file.filename else file_path.suffix.lower()
        
        # 加载文档
        documents = load_document(file_path, file_extension)
        if not documents:
            raise HTTPException(status_code=500, detail="Failed to load document content")
        
        # 分割文档
        from app.rag.splitter import split_documents
        split_docs = split_documents(documents)
        
        if not split_docs:
            raise HTTPException(status_code=500, detail="Failed to split document into chunks")
        
        # TODO: 从向量数据库中删除旧文档并添加新文档（需要实现）
        # 这需要根据具体的向量存储实现来完成
        
        logger.info(f"Successfully updated document: {filename}")
        return JSONResponse(
            status_code=200,
            content={
                "code": 200,
                "message": f"Document {filename} updated successfully",
                "filename": filename
            }
        )
    except HTTPException:
        # 重新抛出已知的HTTP异常
        raise
    except Exception as e:
        logger.error(f"Failed to update document {filename}: {e}")
        raise HTTPException(status_code=500, detail="Failed to update document")


@router.get("/knowledge/vectorized")
async def list_vectorized_documents(limit: int = Query(100, ge=1, le=1000, description="返回文档数量限制")):
    """
    列出所有已向量化的文档
    """
    try:
        # 从向量数据库获取所有文档
        documents = retriever.vector_store.get_all_documents(limit=limit)
        
        return {
            "code": 200,
            "message": "success",
            "data": {
                "documents": documents,
                "total": len(documents)
            }
        }
    except Exception as e:
        logger.error(f"Failed to list vectorized documents: {e}")
        raise HTTPException(status_code=500, detail="Failed to list vectorized documents")
